import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle

# 1. 栅格地图和场景设置
grid_size = 200  # 栅格地图大小 200x200
grid_res = 0.1  # 栅格边长 0.1m
map = np.zeros((grid_size, grid_size))  # 初始化栅格地图 (0表示空间)

# 计算墙体和无人机的坐标（以米为单位，中心为(10,10)）
map_center_m = grid_size * grid_res / 2

# 设置两个墙体 (用1表示被占用)
wall_corner_m = np.array([5, 5])
wall_length_m = 6
wall_corner_idx = np.round(wall_corner_m / grid_res).astype(int)
wall_length_idx = np.round(wall_length_m / grid_res).astype(int)

# 墙体1 (水平部分)
map[wall_corner_idx[1], wall_corner_idx[0]:wall_corner_idx[0] + wall_length_idx] = 1
# 墙体2 (垂直部分)
map[wall_corner_idx[1]:wall_corner_idx[1] + wall_length_idx, wall_corner_idx[0]] = 1


# 设置4架无人机（用1表示被占用）
drones_m = np.array([
    [6.0, 6.0],
    [7.9, 7.9],
    [7.2, 6.0],
    [6.0, 7.2]
])
# 将无人机坐标转换为栅格索引
drones_idx = np.round(drones_m / grid_res).astype(int)
for i in range(drones_idx.shape[0]):
    map[drones_idx[i, 1], drones_idx[i, 0]] = 1

# 2. 预先计算地图信息
warning_zone_radius_m = 1.5
warning_zone_radius_idx = np.round(warning_zone_radius_m / grid_res).astype(int)
nearby_obs_radius_m = 3.0
min_distance_constraint = 0.1
step_distance_m = 0.5

X, Y = np.meshgrid(np.arange(grid_size), np.arange(grid_size))
wall_y, wall_x = np.where(map == 1)
walls_m = np.stack((wall_x, wall_y), axis=-1) * grid_res

# 3. 遍历每个无人机和方向，计算风险
min_risk = float('inf')  # 最小风险，初始化为无穷大
best_drone_idx = -1  # 最佳无人机索引
best_direction_deg = -1  # 最佳方向（角度）
best_next_pos_m = np.array([0, 0])  # 最佳下一步位置 (米)

# 遍历每个无人机
for i in range(drones_m.shape[0]):
    current_drone_m = drones_m[i, :]
    
    # 在此循环内动态创建威胁地图
    threat_map = map.copy()
    other_drones_m = drones_m.copy()
    other_drones_m = np.delete(other_drones_m, i, axis=0)
    
    # 标记其他无人机周围的警戒区域
    for k in range(other_drones_m.shape[0]):
        drone_x = drones_idx[np.where(np.all(drones_m == other_drones_m[k, :], axis=1))[0][0], 0]
        drone_y = drones_idx[np.where(np.all(drones_m == other_drones_m[k, :], axis=1))[0][0], 1]
        distance = np.sqrt((X - drone_x) ** 2 + (Y - drone_y) ** 2)
        threat_map[distance <= warning_zone_radius_idx] = 2

    # 确保墙体和无人机在威胁地图中被正确标记
    threat_map[map == 1] = 1

    # 找到所有威胁的坐标（墙体、其他无人机本身和其预警区域）
    all_threat_y, all_threat_x = np.where(threat_map > 0)
    all_threats_m = np.stack((all_threat_x, all_threat_y), axis=-1) * grid_res

    # 遍历无人机下一步可行位置的方向 (每10度)
    for direction_deg in range(0, 360, 10):
        # 将角度转换为弧度
        direction_rad = np.deg2rad(direction_deg)
        
        # 计算下一步位置的坐标 (米)
        next_pos_m = current_drone_m + np.array([np.cos(direction_rad), np.sin(direction_rad)]) * step_distance_m
        
        # 计算风险
        risk = 0
        
        # 遍历所有威胁
        for j in range(all_threats_m.shape[0]):
            threat_m = all_threats_m[j, :]
            
            # 计算距离，施加最小距离约束
            dist = max(np.linalg.norm(next_pos_m - threat_m), min_distance_constraint)
            
            if dist <= nearby_obs_radius_m:
                # 累加风险 (1/距离^2)
                risk += (1 / dist ** 2)
        
        # 更新最佳路径
        if risk < min_risk:
            min_risk = risk
            best_drone_idx = i
            best_direction_deg = direction_deg
            best_next_pos_m = next_pos_m

# 将最佳下一步位置转换为栅格索引
best_next_pos_idx = np.round(best_next_pos_m / grid_res).astype(int)

# 设置字体为支持中文的字体，例如 SimHei（黑体）
plt.rcParams['font.family'] = 'SimHei'  # 设置为黑体
plt.rcParams['axes.unicode_minus'] = False  # 防止负号显示为方块
# 4. 可视化结果
plt.figure(figsize=(8, 8))
plt.imshow(map, cmap='gray_r', origin='lower', extent=[0, grid_size, 0, grid_size])

# 绘制所有无人机
plt.scatter(drones_idx[:, 0], drones_idx[:, 1], c='red', label='其他无人机', marker='x', s=100)
# 绘制最佳无人机
plt.scatter(drones_idx[best_drone_idx, 0], drones_idx[best_drone_idx, 1], c='green', label='最佳无人机', marker='x', s=100)
# 绘制最佳下一步位置
plt.scatter(best_next_pos_idx[0], best_next_pos_idx[1], c='green', label='下一步位置', marker='o', s=100)
# 绘制飞行路径
plt.plot([drones_idx[best_drone_idx, 0], best_next_pos_idx[0]], 
         [drones_idx[best_drone_idx, 1], best_next_pos_idx[1]], 'g--', label='飞行路径')

# 添加图例
plt.legend(loc='best')

plt.grid(True)
plt.title('无人机最佳脱离路径')
plt.xlabel('X 栅格索引')
plt.ylabel('Y 栅格索引')
plt.show()
